import json
from typing import List, Set, Optional, Dict
from enum import Enum


# 父类 CodeNode
class CodeNode:
    def __init__(self, id: int, name: str, full_qualified_name: str, absolute_path: str = "",
                 start_line: int = 0, end_line: int = 0, content: str = "", comment: str = "", package_name: str = ""):
        # Keep id as string for Neo4j 5.x compatibility (element_id is string format)
        self.id = str(id) if id is not None else ""
        self.name = name
        self.full_qualified_name = full_qualified_name
        self.absolute_path = absolute_path
        self.start_line = start_line
        self.end_line = end_line
        self.content = content
        self.comment = comment
        self.package_name = package_name

    def __repr__(self):
        return f"CodeNode(name='{self.name}', full_qualified_name='{self.full_qualified_name}')"


# Clazz 子类
class Clazz(CodeNode):
    def __init__(self, id: int, name: str, full_qualified_name: str, absolute_path: str = "", start_line: int = 0,
                 end_line: int = 0, content: str = "", comment: str = "", modifiers: Optional[str] = None,
                 class_type: Optional[str] = None, simple_content: Optional[str] = None,
                 parent_classes: Optional[List[str]] = None, summarization: Optional[str] = None, package_name: str = ""):
        super().__init__(id, name, full_qualified_name, absolute_path, start_line, end_line, content, comment, package_name)
        self.modifiers = modifiers
        self.class_type = class_type
        self.simple_content = simple_content
        self.parent_classes = parent_classes or []
        self.summarization = summarization

    def __repr__(self):
        return f"Clazz(name='{self.name}', class_type='{self.class_type}')"


# Method 子类
class Method(CodeNode):
    class Type(Enum):
        METHOD = "METHOD"
        CONSTRUCTOR = "CONSTRUCTOR"
        TESTMETHOD = "TEST_METHOD"

    def __init__(self, id: int, name: str, full_qualified_name: str, absolute_path: str = "", start_line: int = 0,
                 end_line: int = 0, content: str = "", comment: str = "", params: Optional[List[str]] = None,
                 return_type: Optional[str] = None, modifiers: Optional[str] = None,
                 exception_throw: Optional[List[str]] = None, signature: Optional[str] = None,
                 class_name: Optional[str] = None, type: Optional['Method.Type'] = None,
                 summarization: Optional[str] = None, package_name: str = ""):
        super().__init__(id, name, full_qualified_name, absolute_path, start_line, end_line, content, comment, package_name)
        self.params = params or []
        self.return_type = return_type
        self.modifiers = modifiers
        self.exception_throw = exception_throw or []
        self.signature = signature
        self.class_name = class_name
        self.type = type
        self.summarization = summarization

    def __repr__(self):
        return f"Method(name='{self.name}', signature='{self.signature}', type='{self.type}')"

    def to_dict(self):
        return {
            "id": self.id,
            "name": self.name,
            "full_qualified_name": self.full_qualified_name,
            "absolute_path": self.absolute_path,
            "start_line": self.start_line,
            "end_line": self.end_line,
            "content": self.content,
            "comment": self.comment,
            "params": self.params,
            "return_type": self.return_type,
            "modifiers": self.modifiers,
            "exception_throw": self.exception_throw,
            "signature": self.signature,
            "class_name": self.class_name,
            "type": self.type.value if self.type else None,
            "summarization": self.summarization,
            "package_name": self.package_name
        }


# Variable 子类
class Variable(CodeNode):
    def __init__(self, id: int, name: str, full_qualified_name: str, absolute_path: str = "", start_line: int = 0,
                 end_line: int = 0, content: str = "", comment: str = "", modifiers: Optional[str] = None,
                 data_type: Optional[str] = None, signature: Optional[str] = None,
                 class_name: Optional[str] = None, summarization: Optional[str] = None, package_name: str = ""):
        super().__init__(id, name, full_qualified_name, absolute_path, start_line, end_line, content, comment, package_name)
        self.modifiers = modifiers
        self.data_type = data_type
        self.signature = signature
        self.class_name = class_name
        self.summarization = summarization

    def __repr__(self):
        return f"Variable(name='{self.name}', data_type='{self.data_type}')"


# TestClazz 子类
class TestClazz(CodeNode):
    def __init__(self, id: int, name: str, full_qualified_name: str, absolute_path: str = "", start_line: int = 0,
                 end_line: int = 0, content: str = "", comment: str = "", summarization: Optional[str] = None,
                 package_name: str = "", tested_clazz_fqname: Optional[str] = "", tested_method_fqname: Optional[str] = "",
                 method_signature: Optional[str] = "", test_report: Optional[str] = "",
                 coverage_rate: Optional[str] = "", coverage_lines: Optional[List[int]] = None,
                 mutation_score: Optional[str] = "", mutants: Optional[str] = None, find_bugs: Optional[bool] = False,
                 related_nodes: Optional[List[int]] = None, test_points: Optional[str] = "", requirement: Optional[str] = ""):
        super().__init__(id, name, full_qualified_name, absolute_path, start_line, end_line, content, comment, package_name)
        self.summarization = summarization
        self.tested_clazz_FQname = tested_clazz_fqname
        self.tested_method_FQname = tested_method_fqname
        self.method_signature = method_signature
        self.requirement = requirement
        self.test_report = test_report
        self.coverage_rate = coverage_rate
        self.coverage_lines = coverage_lines or []
        self.mutation_score = mutation_score
        self.mutants = json.loads(mutants) if mutants else {}
        self.related_nodes = related_nodes or []
        self.find_bugs = find_bugs
        self.test_points = json.loads(test_points) if test_points else {}

    def __repr__(self):
        return f"Clazz(name='{self.name}', tested_clazz_FQname='{self.tested_clazz_FQname}', focal_method='{self.method_signature})"


def _convert_node(node):
    """
    将 Neo4j 节点转换为对应的类实例，根据节点的标签进行分类。
    分为 Clazz, Method, Variable 三种标签。
    """
    if "Clazz" in node.labels:
        return _convert_to_clazz(node)
    elif "Method" in node.labels:
        return _convert_to_method(node)
    elif "Variable" in node.labels:
        return _convert_to_variable(node)
    elif "TestClazz" in node.labels:
        return _convert_to_testClazz(node)


def _convert_to_clazz(node):
    return Clazz(
        id=node.element_id,
        name=node["name"],
        full_qualified_name=node["full_qualified_name"],
        absolute_path=node.get("absolute_path", ""),
        start_line=node.get("start_line", 0),
        end_line=node.get("end_line", 0),
        content=node.get("content", ""),
        comment=node.get("comment", ""),
        modifiers=node.get("modifiers", []),
        class_type=node.get("class_type", ""),
        simple_content=node.get("simple_content", ""),
        parent_classes=node.get("parent_classes", []),
        summarization=node.get("summarization", ""),
        package_name=node.get("package_name", "")
    )


def _convert_to_method(node):
    return Method(
        id=node.element_id,
        name=node["name"],
        full_qualified_name=node["full_qualified_name"],
        absolute_path=node.get("absolute_path", ""),
        start_line=node.get("start_line", 0),
        end_line=node.get("end_line", 0),
        content=node.get("content", ""),
        comment=node.get("comment", ""),
        params=node.get("params", []),
        return_type=node.get("return_type", ""),
        modifiers=node.get("modifiers", []),
        exception_throw=node.get("exception_throw", []),
        signature=node.get("signature", ""),
        class_name=node.get("class_name", ""),
        type=Method.Type[node.get("type", "METHOD").upper()] if node.get("type") else None,
        summarization=node.get("summarization", ""),
        package_name=node.get("package_name", "")
    )


def _convert_to_variable(node):
    return Variable(
        id=node.element_id,
        name=node["name"],
        full_qualified_name=node["full_qualified_name"],
        absolute_path=node.get("absolute_path", ""),
        start_line=node.get("start_line", 0),
        end_line=node.get("end_line", 0),
        content=node.get("content", ""),
        comment=node.get("comment", ""),
        modifiers=node.get("modifiers", []),
        data_type=node.get("data_type", ""),
        signature=node.get("signature", ""),
        class_name=node.get("class_name", ""),
        summarization=node.get("summarization", ""),
        package_name=node.get("package_name", "")
    )


def _convert_to_testClazz(node):
    return TestClazz(
        id=node.element_id,
        name=node["name"],
        full_qualified_name=node["full_qualified_name"],
        absolute_path=node.get("absolute_path", ""),
        start_line=node.get("start_line", 0),
        end_line=node.get("end_line", 0),
        content=node.get("content", ""),
        comment=node.get("comment", ""),
        summarization=node.get("summarization", ""),
        package_name=node.get("package_name", ""),
        tested_clazz_fqname=node.get("tested_clazz_FQname", ""),
        tested_method_fqname=node.get("tested_method_FQname", ""),
        method_signature=node.get("method_signature", ""),
        requirement=node.get("requirement", ""),
        test_report=node.get("test_report", ""),
        coverage_rate=node.get("coverage_rate", ""),
        coverage_lines=node.get("coverage_lines", []),
        mutation_score=node.get("mutation_score", ""),
        mutants=node.get("mutants", []),
        related_nodes=node.get("related_nodes", [])
    )
